import keras
from PIL import Image
from keras.models import Sequential
from keras.layers import Dense, Dropout, Flatten, Activation
from keras.layers import Conv2D, MaxPooling2D
from keras import backend as k
from keras.optimizers import SGD
from keras.utils import np_utils
import numpy as np
from theano.sparse.basic import le


def load_images(_url):
	# 加载照片
	img = Image.open(_url)
	img_array = np.asarray(img, dtype='float64') / 255
	# 57 x 47 = 2679
	faces = np.empty((400, 2679))
	for row in range(20):
		for col in range(20):
			faces[row * 20 + col] = np.ndarray.flatten(img_array[row * 57: (row + 1) * 57, col * 47])
	label = np.empty(400)
	for i in range(40):
		label[i * 10: i * 10 + 10] = i
	label = label.astype(np.int)

	train_data = np.empty((320, 2679))
	train_label = np.empty(320)
	valid_data = np.empty((40, 279))
	valid_label = np.empty(40)
	test_data = np.empty((40, 2679))
	test_label = np.empty(40)
	for i in range(40):
		train_data[i * 8: i * 8 + 8] = faces[i * 10: i * 10 + 8]
		train_label[i * 8: i * 8 + 8] = label[i * 10: i * 10 + 8]
		valid_data[i] = faces[i * 10 + 8]
		valid_label[i] = label[i * 10 + 8]
		test_data[i] = faces[i * 10 + 9]
		test_label[i] = label[i * 10 + 9]
	train_data = train_data.astype('float32')
	valid_data = valid_data.astype('float32')
	test_data = test_data.astype('float32')
	rval = [(train_data, train_label), (valid_data, valid_label), (test_data, test_label)]
	return rval


def set_model(lr=0.005, decay=le-6, momentum=0.9):
	model = Sequential()
	if k.image_data_format() == 'channel_first':
		model.add(Conv2D(5, kernel_size=(3, 3), input_shape=(1, 57, 47)))
	else:
		model.add(Conv2D(5, kernel_size=(3, 3), input_shape=(57, 47, 1)))
	model.add(Activation('tanh'))
	# model.add(MaxPooling2D(pool_size=(2, 3)))
	model.add(Conv2D(10, kernel_size=(3, 3)))
	model.add(Activation('tanh'))
	model.add(MaxPooling2D(pool_size=(2, 2)))
	model.add(Dropout(0.25))
	model.add(Flatten())
	model.add(Dense(128))
	model.add(Activation('tanh'))
	model.add(Dropout(0.5))
	model.add(Dense(40))
	model.add(Activation('softmax'))
	sgd = SGD(lr=lr, decay=decay, momentum=momentum, nesterov=True)
	model.compile(loss='categorical_crossentropy', optimizer=sgd)
	return model


def train_model(model, x_train, y_train, x_val, y_val):
	model.fit(x_train, y_train, batch_size=40, epochs=epochs, verbose=1, validation_data=(x_val, y_val))
	model.save_weights('model_weights.h5', overwrite=True)
	return model


def test_model(model, x, y):
	model.load_weights('model_weights.h5')
	score = model.evaluate(x, y, verbose=0)
	return score


np.random.seed(1337)
# 类别(大于样本数)
nb_classes = 40
# 学习次数
epochs = 40
# 一次训练的数量
batch_size = 40
# 照片大小
img_row, img_col = 57, 47
# 卷积核的大小
nb_conv = 3
# 池化大小
nb_pool = 2
# 卷积核个数
nb_filters1, nb_filters2 = 5, 10

if __name__ == "__main__":
	img_url = ''
	(x_train, y_train), (x_val, y_val), (x_test, y_test) = load_images(_url=img_url)
	if k.image_data_format() == 'channels_first':
		x_train = x_train.reshape(x_train.shape[0], 1, img_row, img_col)
		x_val = x_val.reshape(x_val.shape[0], 1, img_row, img_col)
		x_test = x_test.reshape(x_test.shape[0], 1, img_row, img_col)
		input_shape = (1, img_row, img_col)
	else:
		x_train = x_train.reshape(x_train.shape[0], img_row, img_col, 1)
		x_val = x_val.reshape(x_val.shape[0], img_row, img_col, 1)
		x_test = x_test.reshape(x_test.shape[0], img_row, img_col, 1)
		input_shape = (img_row, img_col, 1)
	Y_train = np_utils.to_categorical(y_train, 40)
	Y_val = np_utils.to_categorical(y_val, 40)
	Y_test = np_utils.to_categorical(y_test, 40)

	model = set_model()
	model = train_model(model, x_train, Y_train, x_val, Y_val)
	score = test_model(model, x_test, Y_test)
	model.load_weights('model_weights.h5')
	classes = model.predict_classes(x_test, verbose=0)
	test_accuracy = np.mean(np.equal(y_test, classes))
	print('accuracy: ', test_accuracy)
